#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File    :   camera_params_convert.py
@Time    :   2024/10/31 16:24:44
@Author  :   YcZ
@Version :   1.0
@Contact :   yuchao.zhou@outlook.com
@License :   (C)Copyright 2022-2024, YcZ
@Desc    :   
'''


import yaml
import argparse
import os

def insert_(a):
    return "_".join(a)

def load_intrinsic_sample():
    return \
        {
            "header": {
                "seq": 0,
                "stamp": {
                    "secs": 0,
                    "nsecs": 0
                },
                "frame_id": ""
            },
            "height": 0,
            "width": 0,
            "distortion_model": "plumb_bob",
            "D": [],
            "K": [],
            "R": [],
            "P": [],
            "binning_x": 0,
            "binning_y": 0,
            "roi": {
                "x_offset": 0,
                "y_offset": 0,
                "height": 0,
                "width": 0,
                "do_rectify": False
            }
        }

def load_extrinsic_sample():
    return \
        {
            "header": {
                "seq": 0,
                "stamp": {
                    "secs": 0,
                    "nsecs": 0
                },
                "frame_id": ""
            },
            "child_frame_id": "",
            "transform": {
                "rotation": {
                    "x": 0.,
                    "y": 0.,
                    "z": 0.,
                    "w": 0.
                },
                "translation": {
                    "x": 0.,
                    "y": 0.,
                    "z": 0.
                }
            }
        }

def represent_sequence(dumper, data):
    return dumper.represent_sequence(u'tag:yaml.org,2002:seq', data, flow_style=True)

def dump_yaml(filename, data):
    yaml.add_representer(list, represent_sequence)
    with open(filename, 'w') as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=False)

def parse_and_convert_intrinsic(yaml_file, output_folder):
    data = None
    with open(yaml_file, 'r') as f:
        data = yaml.load(f, Loader=yaml.FullLoader)

    filenames = os.path.basename(yaml_file).split('_')
    no = filenames[2]
    pos = insert_(filenames[3:-3])
    focus = insert_(filenames[-3:-1])
    frame_id = insert_([pos, no])
    camera_intrinsics = load_intrinsic_sample()
    camera_intrinsics['header']['frame_id'] = frame_id
    camera_intrinsics['height'] = data['height']
    camera_intrinsics['width'] = data['width']

    camera_intrinsics['K'] = data['K']
    camera_intrinsics['D'] = list(data['D'])
    camera_intrinsics['R'] = data['R']
    camera_intrinsics['P'] = data['P']
    # dump_yaml(os.path.join(output_folder, insert_([frame_id, focus, "intrinsics"])+'.yaml'), camera_intrinsics)
    dump_yaml(os.path.join(output_folder, insert_([frame_id, "intrinsics"])+'.yaml'), camera_intrinsics)

def parse_and_convert_extrinsic(yaml_file, output_folder):
    with open(yaml_file, 'r') as f:
        data = yaml.load(f, Loader=yaml.FullLoader)

    if "to" in yaml_file: # camera to lidar
        pre_filenames = os.path.basename(yaml_file).split('_to_')[0]
        filenames = pre_filenames.split('_')
        no = filenames[2]
        pos = insert_(filenames[3:-2])
        focus = insert_(filenames[-2:])
        child_frame_id = insert_([pos, no])
        frame_id = "lidar32_top"
    else: # camera to novatel
        filenames = os.path.basename(yaml_file).split('_')
        no = filenames[2]
        pos = insert_(filenames[3:-3])
        focus = insert_(filenames[-3:-1])
        child_frame_id = insert_([pos, no])
        frame_id = "novatel"

    camera_extrinsics = load_extrinsic_sample()
    camera_extrinsics['header']['frame_id'] = frame_id
    camera_extrinsics['child_frame_id'] = child_frame_id
    camera_extrinsics['transform'] = data['transform']
    # camera_extrinsics['rotation'] = data['rotation']
    # camera_extrinsics['translation'] = data['translation']
    # dump_yaml(os.path.join(output_folder, insert_([child_frame_id, focus, "to", frame_id, "extrinsics"])+'.yaml'), camera_extrinsics)
    dump_yaml(os.path.join(output_folder, insert_([child_frame_id, "to", frame_id, "extrinsics"])+'.yaml'), camera_extrinsics)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-d','--directory', type=str, default='./orin_camera_params')
    parser.add_argument('-o','--output', type=str, default='./params_output3')
    args = parser.parse_args()

    folder = args.directory
    output_folder = args.output
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    for file in os.listdir(folder):
        print("checking",file)
        if file.endswith('.yml') or file.endswith('.yaml'):
            if 'intrinsics' in file:
                print("convert intrinsic", file)
                parse_and_convert_intrinsic(os.path.join(folder, file), output_folder)
            if 'extrinsics' in file:
                print("convert extrinsic", file)
                parse_and_convert_extrinsic(os.path.join(folder, file), output_folder)
            else:
                print("pass", file)
        else:
            continue
        # input()
    print("Done")

